import cv2
import os
import json
from typing import Tuple
import numpy as np

font = cv2.FONT_HERSHEY_SIMPLEX

# 解析$HOME环境变量，构建配置文件的路径
config_dir = os.path.expanduser('~/.bc_config')
CONFIG_FILE = os.path.join(config_dir, 'cam_calibration.config')

WINDOW_W = 960
WINDOW_H = 720
WINDOW_NAME = "Cam Calibration"

# 目标方块的边长为3.9
CUBE_WIDTH = 3.9

# 摄像头标定工具，包含颜色检测、像素密度计算等功能
class ColorDetector:
    """
    用于检测图像中的颜色阈值
    """
    def __init__(
        self, window_name="Image", hsv_range: Tuple[int, int, int] = (20, 50, 50)
    ):
        self.gray_frame = None
        self.hsv_frame = None
        self.cnts_list = []
        self.window_name = window_name
        self.hsv_range = hsv_range

    """根据鼠标点击调整滑动条的数值。
    此函数通过更新滑动条的数值来检测颜色阈值。
    :param event: 鼠标点击事件，例如左键点击。
    :param x: int 像素坐标
    :param y: int 像素坐标
    :param flags: int
    :param param: NoneType
    """
    def mouse_click(self, event, x, y, flags, param):
        # 鼠标点击事件，更新滑动条的数值
        if event == cv2.EVENT_LBUTTONDOWN:
            self.update_trackbar_values(x, y)

    """更新滑动条的值为点击位置的HSV值。
    :param x: int 像素坐标
    :param y: int 像素坐标
    """
    def update_trackbar_values(self, x, y):
        # 更新滑动条的值为点击位置的HSV值
        h, s, v = self.hsv_frame[y, x]
        cv2.setTrackbarPos("H_MIN", self.window_name, max(0, h - self.hsv_range[0]))
        cv2.setTrackbarPos("H_MAX", self.window_name, min(180, h + self.hsv_range[0]))
        cv2.setTrackbarPos("S_MIN", self.window_name, max(0, s - self.hsv_range[1]))
        cv2.setTrackbarPos("S_MAX", self.window_name, min(255, s + self.hsv_range[1]))
        cv2.setTrackbarPos("V_MIN", self.window_name, max(0, v - self.hsv_range[2]))
        cv2.setTrackbarPos("V_MAX", self.window_name, min(255, v + self.hsv_range[2]))

    """显示像素信息。
    :param x: int 像素坐标
    :param y: int 像素坐标
    """
    def display_pixel_info(self, x, y):
        # 显示像素信息
        print("PIX: ", x, y)
        print("GRAY: ", self.gray_frame[y, x])
        print("HSV: ", self.hsv_frame[y, x])

    """检测指定阈值内的颜色块。
    此函数用矩形框突出显示指定颜色阈值内的块。
    :param img: ndarray, 视频帧。
    :param h_min: int, 最小色调
    :param h_max: int, 最大色调
    :param s_min: int, 最小饱和度
    :param s_max: int, 最大饱和度
    :param v_min: int, 最小明度
    :param v_max: int, 最大明度
    """
    def update_frame(self, img, h_min, h_max, s_min, s_max, v_min, v_max):
        # 根据阈值更新图像，返回轮廓和处理后图像
        self.convert_frames(img)
        mask_color = self.apply_color_mask(h_min, h_max, s_min, s_max, v_min, v_max)
        self.display_mask(mask_color)
        cnts, frame = self.find_contours(img, mask_color)
        return cnts, frame

    """将图像帧转换为灰度和HSV格式。"""
    def convert_frames(self, img):
        # 将图像转换为灰度和HSV格式
        self.gray_frame = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        self.hsv_frame = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    """应用颜色掩码到HSV帧。"""
    def apply_color_mask(self, h_min, h_max, s_min, s_max, v_min, v_max):
        # 应用颜色掩码到HSV图像
        low_color = np.array([h_min, s_min, v_min])
        high_color = np.array([h_max, s_max, v_max])
        mask_color = cv2.inRange(self.hsv_frame, low_color, high_color)
        return cv2.medianBlur(mask_color, 7)
        
    """显示颜色掩码。"""
    def display_mask(self, mask_color):
        # 显示颜色掩码
        cv2.imshow("mask", mask_color)

    """提取联通阈。 """
    def find_contours(self, img, mask_color):
        # 查找并绘制轮廓
        cnts, _ = cv2.findContours(mask_color, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if cnts:
            self.cnts_list = sorted(
                [
                    (int(x + w / 2), int(y + h / 2), int(w), int(h))
                    for (x, y, w, h) in [cv2.boundingRect(cnt) for cnt in cnts]
                    if 30 < w < 70 and 30 < h < 70
                ],
                key=lambda x: x[0],
            )
            for id, (center_x, center_y, w, h) in enumerate(self.cnts_list, start=1):
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2)
                cv2.putText(
                    img,
                    f"target: {id}",
                    (x, y - 5),
                    font,
                    0.7,
                    (0, 0, 255),
                    2,
                )
            return self.cnts_list, img
        else:
            return [], img


def nothing(x):
    # 占位函数
    pass

def run():
    # 主循环，进行摄像头标定
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("无法打开摄像头，请检查连接。")
        return
    else:
        print("摄像头成功打开。")
        
    color_detector = ColorDetector(window_name=WINDOW_NAME)

    cv2.namedWindow(WINDOW_NAME,cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME,WINDOW_W,WINDOW_H)

    cv2.namedWindow(WINDOW_NAME)
    cv2.createTrackbar("H_MIN", WINDOW_NAME, 35, 180, nothing)
    cv2.createTrackbar("H_MAX", WINDOW_NAME, 40, 180, nothing)
    cv2.createTrackbar("S_MIN", WINDOW_NAME, 100, 255, nothing)
    cv2.createTrackbar("S_MAX", WINDOW_NAME, 115, 255, nothing)
    cv2.createTrackbar("V_MIN", WINDOW_NAME, 180, 255, nothing)
    cv2.createTrackbar("V_MAX", WINDOW_NAME, 190, 255, nothing)
    cv2.createTrackbar("SIZE_MIN", WINDOW_NAME, 35, 60, nothing)
    cv2.createTrackbar("SIZE_MAX", WINDOW_NAME, 55, 120, nothing)
    cv2.setMouseCallback(WINDOW_NAME, color_detector.mouse_click)

    print("按'q'关闭摄像头")
    
    while True:
        ret, frame = cap.read()
        
        h_min = cv2.getTrackbarPos("H_MIN", WINDOW_NAME)
        h_max = cv2.getTrackbarPos("H_MAX", WINDOW_NAME)
        s_min = cv2.getTrackbarPos("S_MIN", WINDOW_NAME)
        s_max = cv2.getTrackbarPos("S_MAX", WINDOW_NAME)
        v_min = cv2.getTrackbarPos("V_MIN", WINDOW_NAME)
        v_max = cv2.getTrackbarPos("V_MAX", WINDOW_NAME)

        centers,frame = color_detector.update_frame(frame, h_min, h_max, s_min, s_max, v_min, v_max)

        if len(centers) == 1:
            center_x, center_y, box_lenght_x, box_length_y = centers[0]
            if abs(box_lenght_x - box_length_y) >= 10:
                print("请调整摄像头角度, 方块位置")
            mean_len = (box_lenght_x + box_length_y) / 2
            pixel_per_cm_ratio = int(10 * mean_len / CUBE_WIDTH) / 10
            center_y_adjuested = int(center_y - 19 * pixel_per_cm_ratio+15)
            print(f"\r机械臂基座像素位置为: x = {center_x-10}, y = {center_y_adjuested+10},当前高度下图像像素密度为:{pixel_per_cm_ratio} pixel/cm.按's'键保存.", end="")
        elif len(centers) > 1:
            print("\r检测到多个同颜色校准块, 请确保当前视场内只有一个同颜色校准块", end="")

        if cv2.waitKey(10) & 0xFF == ord('s'):
            if len(centers)>0:
                res = listen_for_key_and_save((center_x, center_y_adjuested, pixel_per_cm_ratio))
                if res:
                    break
                
        cv2.imshow(WINDOW_NAME,frame)


def listen_for_key_and_save(coordinates):
    if coordinates[0] is not None:
        print(f"\n按's'建保存校准坐标: {int(coordinates[0]-10)}, {int(coordinates[1]+10)} 和 像素密度: {coordinates[2]}, 按其它键取消保存")
        key = cv2.waitKey(0)
        if key == ord('s'):
            # 确保配置文件所在目录存在
            cfg = {
                "arm_base_coordinate": 
                    {
                        "x": coordinates[0]-10,
                        "y": coordinates[1]+10
                    },
                "pixel_per_cm_ratio": coordinates[2]
            }
            save_config(cfg, CONFIG_FILE)
            print(f"校准坐标已保存到: {CONFIG_FILE}")
            return True
        else:
            return False

"""
将配置信息保存为 JSON 格式的文件。
输入：
    config (dict): 包含配置信息的字典。
    file (str): 要保存的文件路径。
输出:
    None
"""
def save_config(config, file):
    """
    将配置信息保存为 JSON 格式的文件。
    输入：
        config (dict): 包含配置信息的字典。
        file (str): 要保存的文件路径。
    输出：
        None
    """
    # 获取file的目录部分
    directory = os.path.dirname(file)
    os.makedirs(directory, exist_ok=True)
    with open(file, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=4)

    
if __name__ == '__main__':
    run()
